package com.examples.es;

import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;
import org.apache.flink.streaming.connectors.elasticsearch.ElasticsearchSinkFunction;
import org.apache.flink.streaming.connectors.elasticsearch.RequestIndexer;
import org.apache.flink.streaming.connectors.elasticsearch7.ElasticsearchSink;
import org.apache.http.HttpHost;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.client.Requests;
import org.elasticsearch.search.sort.SortOrder;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class FlinkESSinkTest {
    public static void main(String[] args) throws Exception {
        //构建Flink环境对象
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        //Source：数据的输入
        DataStreamSource<String> source = env.socketTextStream("localhost", 7777);
        //使用ESBuilder构建输出 increment

        HttpHost httpHost = new HttpHost("192.168.0.117", 9200, "http");
        List<HttpHost> httpHosts = new ArrayList<HttpHost>(){{
            add(httpHost);
        }};
        //、查询总数量
        DataStreamSource<Integer> integerESDataStreamSource = env.addSource(new IndicesCountSource(httpHost, "my-index"));
        integerESDataStreamSource.addSink(new SinkFunction<Integer>() {
            @Override
            public void invoke(Integer value, Context context) throws Exception {
                SinkFunction.super.invoke(value, context);
                System.out.println("integerESDataStreamSource value = " + value);
            }
        });
        //、查询索引全部数据
        DataStreamSource<List<HashMap<Object, Object>>> integerESDataStreamSourceALL = env.addSource(new IndicesQueryAllSource(httpHost, "tiktok_posts_comments"));
        integerESDataStreamSourceALL.addSink(new SinkFunction<List<HashMap<Object, Object>>>() {
            @Override
            public void invoke(List<HashMap<Object, Object>> value, Context context) throws Exception {
                SinkFunction.super.invoke(value, context);
                System.out.println("integerESDataStreamSourceALL value = " + value);
            }
        });
        //、增量提取索引全部数据
        DataStreamSource<HashMap<Object, Object>> hashMapDataStreamIncrementSource = env.addSource(new IndicesQueryIncrementSource(httpHost, "tiktok_posts_comments", 100));
        hashMapDataStreamIncrementSource.addSink(new SinkFunction<HashMap<Object, Object>>() {
            @Override
            public void invoke(HashMap<Object, Object> value, Context context) throws Exception {
                SinkFunction.super.invoke(value, context);
                System.out.println("hashMapDataStreamIncrementSource value = " + value);

            }
        });


        //、 批量插入或更新
        ElasticsearchSink.Builder<String> esSinkBuilder = new ElasticsearchSink.Builder<>(httpHosts,
                new ElasticsearchSinkFunction<String>() {
                    public IndexRequest createIndexRequest(String element) {
                        Map<String, String> json = new HashMap<>();
                        json.put("data", element);
                        return Requests.indexRequest()
                                .index("my-index")
                                .id("121222") // 指定了id之后存在则更新，不存在则插入
                                //.type("my-type")
                                .source(json);
                    }
                    @Override
                    public void process(String element, RuntimeContext ctx, RequestIndexer indexer) {
                        indexer.add(createIndexRequest(element));
                    }
                }
        );
        //、 批量删除
        ElasticsearchSink.Builder<String> esSinkBuilder1 = new ElasticsearchSink.Builder<>(httpHosts,
                new ElasticsearchSinkFunction<String>() {
                    @Override
                    public void process(String s, RuntimeContext runtimeContext, RequestIndexer requestIndexer) {
                        //DeleteIndexRequest deleteIndexRequest = Requests.deleteIndexRequest("my-index");
                        DeleteRequest deleteRequest = new DeleteRequest("my-index", s);
                        requestIndexer.add(deleteRequest);
                    }
                }
        );



        esSinkBuilder.setBulkFlushMaxActions(1);
        esSinkBuilder1.setBulkFlushMaxActions(1);
        //Sink：数据的输出
        source.addSink(esSinkBuilder.build());
        source.addSink(esSinkBuilder1.build());
        //执行操作
        env.execute("flink-es");
    }
}
